#ifndef CUDA_VOLUME_RENDERER_H
#define CUDA_VOLUME_RENDERER_H

#include <cutil_inline.h>
#include <cutil_gl_inline.h>
#include <vector_types.h>
#include <vector_functions.h>
#include <driver_functions.h>

#ifndef CUDA_SAFE_FREE_ARRAY
#define CUDA_SAFE_FREE_ARRAY(ptr) { if(ptr) { cutilSafeCall(cudaFreeArray(ptr)); ptr = NULL; } }
#endif

template<typename T>
void copyVolume(int nx, int ny, int nz, int pitch, T *dst, T *src);
void renderVolume(int nx, int ny, int nz, const cudaArray *volumeArray, const cudaChannelFormatDesc &channelDesc, uint width, uint height, uint *image);

template<typename T>
class VolumeData
{
public:
	cudaExtent volumeSize;
	cudaArray *volumeArray;
	cudaChannelFormatDesc channelDesc;

public:
	VolumeData()
	{
		volumeSize = make_cudaExtent(0, 0, 0);
		volumeArray = NULL;
	}

public:
	void resize(int nx, int ny, int nz)
	{
		CUDA_SAFE_FREE_ARRAY(volumeArray);
		cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(sizeof(T) * 8, 0, 0, 0, cudaChannelFormatKindFloat);
		volumeSize = make_cudaExtent(nx, ny, nz);
		cutilSafeCall(cudaMalloc3DArray(&volumeArray, &channelDesc, volumeSize));
	}

	void update(int nx, int ny, int nz, T *ptr)
	{
		resize(nx, ny, nz);
		cudaPitchedPtr pptr;
		cutilSafeCall(cudaMalloc3D(&pptr, volumeSize));
		copyVolume(nx, ny, nz, pptr.pitch, (T *)pptr.ptr, ptr);

		cudaMemcpy3DParms copyParams = { 0 };
		copyParams.extent = volumeSize;
		copyParams.dstArray = volumeArray;
		copyParams.srcPtr = pptr;
		copyParams.kind = cudaMemcpyDeviceToDevice;
		cutilSafeCall(cudaMemcpy3D(&copyParams));
		cutilSafeCall(cudaFree(pptr.ptr));
	}
};

class VolumeRenderer
{
public:
	int width, height;
	GLuint pbo;
	GLuint tex;
	cudaGraphicsResource *pboRes;

public:
	VolumeRenderer()
	{
		width = 0;
		height = 0;
		pbo = 0;
		tex = 0;
		pboRes = NULL;

		static bool setGLDevice = false;
		if(!setGLDevice)
		{
			cudaGLSetGLDevice(cutGetMaxGflopsDeviceId());
			setGLDevice = true;
		}
	}

	~VolumeRenderer()
	{
		free();
	}

public:
	void free()
	{
		if(pbo)
		{
			cutilSafeCall(cudaGraphicsUnregisterResource(pboRes));
			glDeleteBuffers(1, &pbo);
			pbo = 0;
		}

		if(tex)
		{
			glDeleteTextures(1, &tex);
			tex = 0;
		}
	}

	void resize(int width, int height)
	{
		free();
		this->width = width;
		this->height = height;

		glGenBuffers(1, &pbo);
		glBindBuffer(GL_PIXEL_UNPACK_BUFFER_ARB, pbo);
		glBufferData(GL_PIXEL_UNPACK_BUFFER_ARB, width * height * sizeof(GLubyte) * 4, 0, GL_STREAM_DRAW);
		glBindBuffer(GL_PIXEL_UNPACK_BUFFER_ARB, 0);
		cutilSafeCall(cudaGraphicsGLRegisterBuffer(&pboRes, pbo, cudaGraphicsMapFlagsWriteDiscard));

		glGenTextures(1, &tex);
		glBindTexture(GL_TEXTURE_2D, tex);
		glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, NULL);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
		glBindTexture(GL_TEXTURE_2D, 0);
	}

	void render(const VolumeData &vol)
	{
		uint *image;
		size_t bytes;
		cutilSafeCall(cudaGraphicsMapResources(1, &pboRes, 0));
		cutilSafeCall(cudaGraphicsResourceGetMappedPointer((void **)&image, &bytes, pboRes));

		cutilSafeCall(cudaMemset(image, 0, width * height * sizeof(uint)));
		renderVolume(vol.volumeSize.width, vol.volumeSize.height, vol.volumeSize.depth, vol.volumeArray, vol.channelDesc, width, height, image);
		cutilSafeCall(cudaGraphicsUnmapResources(1, &pboRes, 0));

		glBindBuffer(GL_PIXEL_UNPACK_BUFFER, pbo);
		glBindTexture(GL_TEXTURE_2D, tex);
		glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE, 0);
		glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);
	}
};

#endif
